import os
import shutil
from tqdm import tqdm

def merge_yolo_datasets(det_path, seg_path, merged_path):
    sets = ['train', 'val', 'test']
    
    for s in sets:
        img_dir_det = os.path.join(det_path, 'images', s)
        lbl_dir_det = os.path.join(det_path, 'labels', s)
        img_dir_seg = os.path.join(seg_path, 'images', s)
        lbl_dir_seg = os.path.join(seg_path, 'labels', s)

        merged_img_dir = os.path.join(merged_path, 'images', s)
        merged_lbl_dir = os.path.join(merged_path, 'labels', s)

        os.makedirs(merged_img_dir, exist_ok=True)
        os.makedirs(merged_lbl_dir, exist_ok=True)

        # Merge detection images and labels
        for f in tqdm(os.listdir(img_dir_det), desc=f'Merging det images ({s})'):
            src = os.path.join(img_dir_det, f)
            dst = os.path.join(merged_img_dir, f)
            shutil.copy2(src, dst)
        
        for f in tqdm(os.listdir(lbl_dir_det), desc=f'Merging det labels ({s})'):
            src = os.path.join(lbl_dir_det, f)
            dst = os.path.join(merged_lbl_dir, f)
            shutil.copy2(src, dst)

        # Merge segmentation images and labels
        for f in tqdm(os.listdir(img_dir_seg), desc=f'Merging seg images ({s})'):
            src = os.path.join(img_dir_seg, f)
            dst = os.path.join(merged_img_dir, f)

            # Prevent overwriting by renaming if needed
            if os.path.exists(dst):
                base, ext = os.path.splitext(f)
                dst = os.path.join(merged_img_dir, base + '_seg' + ext)
            shutil.copy2(src, dst)

        for f in tqdm(os.listdir(lbl_dir_seg), desc=f'Merging seg labels ({s})'):
            src = os.path.join(lbl_dir_seg, f)
            dst = os.path.join(merged_lbl_dir, f)

            # Prevent overwriting by renaming if needed
            if os.path.exists(dst):
                base, ext = os.path.splitext(f)
                dst = os.path.join(merged_lbl_dir, base + '_seg' + ext)
            shutil.copy2(src, dst)

if __name__ == "__main__":
    # det_path = "/home/lhx/pcbsyn/data_zoo/splittruecolorseg"
    # seg_path = "/home/lhx/pcbsyn/data_zoo/splitsyntruecolor"
    det_path="/home/lhx/pcbsyn/data_zoo/splittruecolorwithdepthseg"
    seg_path="/home/lhx/pcbsyn/data_zoo/splitsynpcbwithjrsdepbg"
    merged_path = "/home/lhx/pcbsyn/data_zoo/mergetruecolordep"
    merge_yolo_datasets(det_path, seg_path, merged_path)
